import nltk
from nltk.corpus import brown
from nltk import FreqDist
# import matplotlib.pyplot as plt

# 下载语料库（首次运行需要）
nltk.download('brown', quiet=True)

# 获取语料库单词并转换为小写
words = [word.lower() for word in brown.words()]

# 计算词频分布
freq_dist = FreqDist(words)
print(f"语料库总词汇量: {len(freq_dist)}")

# 获取按频率降序排列的所有词汇
sorted_word_freq = freq_dist.most_common()

# 确定目标词汇范围
start_rank = 7000  # 排名7001开始（索引7000）
end_rank = 10000   # 排名10000结束（索引9999）

# 提取7000-10000排名的词汇
target_words = [word for word, freq in sorted_word_freq[start_rank:end_rank]]

# 创建词汇表
print("\n词汇表 (排名7001-10000):")
print(f"共 {len(target_words)} 个词汇\n")

# 分组打印词汇（每行10个）
for i in range(0, len(target_words), 10):
    print(", ".join(target_words[i:i+10]))

# 可视化词频分布（可选）
# plt.figure(figsize=(12, 6))
frequencies = [freq for _, freq in sorted_word_freq[start_rank:end_rank]]
# plt.plot(range(start_rank, end_rank), frequencies)
# plt.title(f'词汇频率分布 (排名 {start_rank+1}-{end_rank})')
# plt.xlabel('词汇排名')
# plt.ylabel('出现频率')
# plt.grid(True)
# plt.savefig('vocabulary_frequency.png')
# plt.show()

# 保存到文件（可选）
with open('vocabulary_expansion.txt', 'w') as f:
    f.write("\n".join(target_words))
print("\n词汇表已保存到 vocabulary_expansion.txt")